import gzip
from abc import ABC, abstractmethod
from enum import IntEnum
from typing import Optional


class CodecType(IntEnum):
    GZIP = 0
    LZ4 = 1
    ZSTD = 2
    ZIP = 3
    BZ2 = 4

    @staticmethod
    def get_codec_type_by_name(name: str) -> Optional[IntEnum]:
        for codec in CodecType:
            if codec.name == name.upper():
                return codec
        return None

    @staticmethod
    def get_codec_type_by_bitwise(bitwise: int) -> Optional[IntEnum]:
        for codec in CodecType:
            if codec.value == bitwise:
                return codec
        return None


class Compressor(ABC):
    @abstractmethod
    def compress(self, bytes_: bytes) -> bytes:
        pass


class Decompressor(ABC):
    @abstractmethod
    def decompress(self, bytes_: bytes) -> bytes:
        pass


class Codec(Compressor, Decompressor):
    # 标识codec类型，占一个字节
    HEADER_CODEC_TYPE_BYTES = 1
    # 标识压缩后的数据长度，占4个字节，意思支持压缩后的数据最大长度为4GB
    HEADER_CODEC_DATA_LENGTH_BYTES = 4
    # 字节序
    HEADER_CODEC_BYTE_ORDER = 'little'
    # 非符号数
    HEADER_CODEC_SIGNED = False

    @staticmethod
    def get_header_length():
        return Codec.HEADER_CODEC_TYPE_BYTES + Codec.HEADER_CODEC_DATA_LENGTH_BYTES

    @abstractmethod
    def get_codec_type(self):
        pass

    def _generate_header(self, length: int) -> bytes:
        codec_type_int = self.get_codec_type().value
        return codec_type_int.to_bytes(length=Codec.HEADER_CODEC_TYPE_BYTES, byteorder=Codec.HEADER_CODEC_BYTE_ORDER,
                                       signed=Codec.HEADER_CODEC_SIGNED) \
               + length.to_bytes(length=Codec.HEADER_CODEC_DATA_LENGTH_BYTES, byteorder=Codec.HEADER_CODEC_BYTE_ORDER,
                                 signed=Codec.HEADER_CODEC_SIGNED)

    @staticmethod
    def __trim_and_get_header(bytes_: bytes):
        header_len = Codec.HEADER_CODEC_TYPE_BYTES + Codec.HEADER_CODEC_DATA_LENGTH_BYTES
        if len(bytes_) < Codec.get_header_length():
            raise Exception("The format of compressed bytes data is invalid, its header length is less than %s."
                            % Codec.get_header_length())
        return bytes_[:header_len]

    @staticmethod
    def _get_codec_type_from_compressed_data(bytes_: bytes):
        header_bytes = Codec.__trim_and_get_header(bytes_)
        bitwise_bytes = header_bytes[:Codec.HEADER_CODEC_TYPE_BYTES]
        bitwise = int.from_bytes(bytes=bitwise_bytes, byteorder=Codec.HEADER_CODEC_BYTE_ORDER,
                                 signed=Codec.HEADER_CODEC_SIGNED)
        return CodecType.get_codec_type_by_bitwise(bitwise)

    @staticmethod
    def _get_data_length_from_compressed_data(bytes_: bytes):
        header_bytes = Codec.__trim_and_get_header(bytes_)
        length_bytes = header_bytes[Codec.HEADER_CODEC_TYPE_BYTES: Codec.get_header_length()]
        return int.from_bytes(bytes=length_bytes, byteorder=Codec.HEADER_CODEC_BYTE_ORDER, signed=Codec.HEADER_CODEC_SIGNED)


class GzipCodec(Codec):

    def get_codec_type(self):
        return CodecType.GZIP

    def compress(self, bytes_: bytes) -> bytes:
        compressed = gzip.compress(bytes_)
        return self._generate_header(len(compressed)) + compressed

    def decompress(self, bytes_: bytes) -> bytes:
        data_len = self._get_data_length_from_compressed_data(bytes_)
        bytes_ = bytes_[Codec.get_header_length():]
        if len(bytes_) < data_len:
            raise RuntimeError("The data error, the length of valid loaded compressed bytes data is less than "
                               "its length recorded, expect %s but %s" % (data_len, len(bytes_)))
        return gzip.decompress(bytes_[:data_len])


_CODEC_GENERATOR = MAP = {
    CodecType.GZIP: GzipCodec
}


class AutoDecompressor(Codec):
    def get_codec_type(self):
        raise RuntimeError("An unsupported operation.")

    def compress(self, bytes_: bytes) -> bytes:
        raise RuntimeError("An unsupported operation.")

    def __init__(self):
        self._codec_dict = {}

    def decompress(self, bytes_: bytes) -> bytes:
        codec_type = self._get_codec_type_from_compressed_data(bytes_)
        if self._codec_dict.keys().__contains__(codec_type):
            return self._codec_dict[codec_type].decompress(bytes_)
        codec_generator = _CODEC_GENERATOR[codec_type]
        if not codec_generator:
            raise RuntimeError("There not exists codec type %s." % codec_type)
        codec = codec_generator()
        self._codec_dict[codec_type] = codec
        return codec.decompress(bytes_)
